package com.sky;

import java.awt.Color;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import javax.imageio.ImageIO;
import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.awt.BasicStroke;

public class ImageProcessor {
    private Color bboxColor = Color.RED; // 红色边界框
    private int bboxThickness = 3; // 边界框线条粗细

    /**
     * 解析边界框字符串格式 <bbox>x_min y_min x_max y_max</bbox>
     *
     * @param bboxStr 边界框字符串
     * @return 包含四个坐标值的整数数组，或null（解析失败时）
     */
    public int[] parseBboxString(String bboxStr) {
        try {
            if (bboxStr.startsWith("<bbox>") && bboxStr.endsWith("</bbox>")) {
                String coordsStr = bboxStr.substring(6, bboxStr.length() - 7);
                String[] coordsArray = coordsStr.trim().split("\\s+");

                if (coordsArray.length == 4) {
                    int[] coords = new int[4];
                    for (int i = 0; i < 4; i++) {
                        coords[i] = Integer.parseInt(coordsArray[i]);
                    }
                    return coords;
                }
            }
            System.out.println("解析边界框字符串失败: " + bboxStr);
            return null;
        } catch (Exception e) {
            System.out.println("解析边界框字符串时发生异常: " + e.getMessage());
            return null;
        }
    }

    /**
     * 将标准化坐标(0-1000)转换为实际图像坐标
     *
     * @param bbox 标准化坐标数组 [x_min, y_min, x_max, y_max]
     * @param imageWidth 图像宽度
     * @param imageHeight 图像高度
     * @return 实际坐标数组
     */
    public int[] normalizeCoordinates(int[] bbox, int imageWidth, int imageHeight) {
        int xMin = bbox[0];
        int yMin = bbox[1];
        int xMax = bbox[2];
        int yMax = bbox[3];

        int xMinReal = (int) (xMin * imageWidth / 1000.0);
        int yMinReal = (int) (yMin * imageHeight / 1000.0);
        int xMaxReal = (int) (xMax * imageWidth / 1000.0);
        int yMaxReal = (int) (yMax * imageHeight / 1000.0);

        return new int[]{xMinReal, yMinReal, xMaxReal, yMaxReal};
    }

    /**
     * 在图像上绘制边界框
     *
     * @param imageData 图像数据字节数组
     * @param bboxes 边界框列表，每个边界框是包含四个坐标值的整数数组
     * @return 处理后的图像数据字节数组
     * @throws IOException 处理图像时发生的IO异常
     */
    public byte[] drawBboxesOnImage(byte[] imageData, List<int[]> bboxes) throws IOException {
        // 将字节数组转换为BufferedImage
        BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData));

        if (image == null) {
            throw new IllegalArgumentException("无法解码图像数据");
        }

        int width = image.getWidth();
        int height = image.getHeight();
        System.out.println("图像尺寸: " + width + "x" + height);

        // 创建Graphics2D对象进行绘制
        BufferedImage outputImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        Graphics2D g2d = outputImage.createGraphics();

        // 设置抗锯齿，使线条更平滑
        g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        g2d.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);

        // 绘制原始图像
        g2d.drawImage(image, 0, 0, null);

        // 设置边界框绘制属性
        g2d.setColor(bboxColor);
        g2d.setStroke(new BasicStroke(bboxThickness));

        // 在图像上绘制所有边界框
        for (int[] bbox : bboxes) {
            int[] realCoords = normalizeCoordinates(bbox, width, height);

            // 确保坐标在图像范围内
            int xMinReal = Math.max(0, Math.min(realCoords[0], width));
            int yMinReal = Math.max(0, Math.min(realCoords[1], height));
            int xMaxReal = Math.max(0, Math.min(realCoords[2], width));
            int yMaxReal = Math.max(0, Math.min(realCoords[3], height));

            // 绘制矩形边界框
            g2d.drawRect(xMinReal, yMinReal, xMaxReal - xMinReal, yMaxReal - yMinReal);

            System.out.println("绘制边界框: (" + xMinReal + ", " + yMinReal + ", " +
                    xMaxReal + ", " + yMaxReal + ")");
        }

        // 释放资源
        g2d.dispose();

        // 将处理后的图像转换为字节数组
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ImageIO.write(outputImage, "png", baos);
        return baos.toByteArray();
    }

    /**
     * 使用边界框字符串列表处理图像
     *
     * @param imageData 图像数据字节数组
     * @param bboxStrings 边界框字符串列表
     * @return 包含处理后图像数据和有效边界框数量的对象
     */
    public ProcessingResult processImageWithBboxStrings(byte[] imageData, List<String> bboxStrings) {
        List<int[]> bboxes = new ArrayList<>();

        // 解析所有边界框字符串
        for (String bboxStr : bboxStrings) {
            int[] bbox = parseBboxString(bboxStr.trim());
            if (bbox != null) {
                bboxes.add(bbox);
            } else {
                System.out.println("无效的边界框字符串: " + bboxStr);
            }
        }

        // 如果没有有效边界框，返回原始图像
        if (bboxes.isEmpty()) {
            System.out.println("没有有效的边界框，返回原图像");
            return new ProcessingResult(imageData, 0);
        }

        try {
            // 绘制边界框
            byte[] processedImage = drawBboxesOnImage(imageData, bboxes);
            return new ProcessingResult(processedImage, bboxes.size());
        } catch (IOException e) {
            System.out.println("绘制边界框时发生IO异常: " + e.getMessage());
            return new ProcessingResult(imageData, 0);
        }
    }

    // 用于返回处理结果的内部类
    public static class ProcessingResult {
        private final byte[] imageData;
        private final int validBboxCount;

        public ProcessingResult(byte[] imageData, int validBboxCount) {
            this.imageData = imageData;
            this.validBboxCount = validBboxCount;
        }

        public byte[] getImageData() {
            return imageData;
        }

        public int getValidBboxCount() {
            return validBboxCount;
        }
    }

    // 设置边界框颜色
    public void setBboxColor(Color color) {
        this.bboxColor = color;
    }

    // 设置边界框线条粗细
    public void setBboxThickness(int thickness) {
        this.bboxThickness = thickness;
    }
}
